import { WebSocket } from 'ws';
import { Endpoint, EndpointStatus } from '../endpoint';
import { RemoteConfig } from '../remote';
import { RemoteMessage, RemoteMessageType } from '../remoteProcess';

export class WebSocketEndpoint extends Endpoint {
  private ws?: WebSocket;

  constructor(address: string, config: RemoteConfig) {
    super(address, config);
  }

  protected async establishConnection(): Promise<void> {
    return new Promise((resolve, reject) => {
      try {
        const url = this.getWebSocketUrl();
        this.ws = new WebSocket(url);

        this.ws.on('open', () => {
          resolve();
        });

        this.ws.on('message', (data: Buffer) => {
          this.handleMessage(data);
        });

        this.ws.on('close', () => {
          this.handleClose();
        });

        this.ws.on('error', (error: Error) => {
          this.handleError(error);
        });
      } catch (error) {
        reject(error);
      }
    });
  }

  async send(message: any): Promise<void> {
    if (!this.ws || this.ws.readyState !== WebSocket.OPEN) {
      throw new Error('WebSocket not connected');
    }

    const serializedMessage = this.serializeMessage(message);
    return new Promise((resolve, reject) => {
      this.ws!.send(serializedMessage, (error) => {
        if (error) {
          reject(error);
        } else {
          resolve();
        }
      });
    });
  }

  protected async disconnect(): Promise<void> {
    if (this.ws) {
      this.ws.close();
      this.ws = undefined;
    }
  }

  private handleMessage(data: Buffer): void {
    try {
      const message = this.deserializeMessage(data);
      if (message.type === RemoteMessageType.Heartbeat) {
        // 处理心跳消息
        this._lastActivity = Date.now();
      } else {
        // 处理其他消息
        this.emit('message', message);
      }
    } catch (error) {
      this.emit('error', error);
    }
  }

  private handleClose(): void {
    this.ws = undefined;
    if (this.status !== EndpointStatus.Closed) {
      this.setStatus(EndpointStatus.Disconnected);
      this.emit('disconnected');
      this.scheduleReconnect();
    }
  }

  private handleError(error: Error): void {
    this.emit('error', error);
    if (this.ws) {
      this.ws.close();
      this.ws = undefined;
    }
  }

  private getWebSocketUrl(): string {
    const protocol = 'ws://';  // TODO: 支持 wss://
    return `${protocol}${this.address}`;
  }

  private serializeMessage(message: any): Buffer {
    if (!this.config.serializer) {
      throw new Error('No serializer configured');
    }
    return Buffer.from(this.config.serializer.serialize(message));
  }

  private deserializeMessage(data: Buffer): RemoteMessage {
    if (!this.config.serializer) {
      throw new Error('No serializer configured');
    }
    return this.config.serializer.deserialize(data, 'RemoteMessage') as RemoteMessage;
  }

  protected async sendHeartbeat(): Promise<void> {
    const heartbeat: RemoteMessage = {
      type: RemoteMessageType.Heartbeat,
      message: {
        timestamp: Date.now()
      }
    };
    await this.send(heartbeat);
  }
} 